package com.ricky.constencrypt.plugin

import com.android.build.api.transform.DirectoryInput
import com.android.build.api.transform.Format
import com.android.build.api.transform.JarInput
import com.android.build.api.transform.QualifiedContent
import com.android.build.api.transform.Status.ADDED
import com.android.build.api.transform.Status.CHANGED
import com.android.build.api.transform.Status.NOTCHANGED
import com.android.build.api.transform.Status.REMOVED
import com.android.build.api.transform.TransformInvocation
import org.apache.commons.compress.archivers.jar.JarArchiveEntry
import org.apache.commons.compress.archivers.zip.ParallelScatterZipCreator
import org.apache.commons.compress.archivers.zip.ZipArchiveEntry
import org.apache.commons.compress.archivers.zip.ZipArchiveOutputStream
import org.apache.commons.compress.parallel.InputStreamSupplier
import org.objectweb.asm.ClassReader
import org.objectweb.asm.ClassVisitor
import org.objectweb.asm.ClassWriter
import org.objectweb.asm.MethodVisitor
import org.objectweb.asm.Opcodes
import org.objectweb.asm.tree.ClassNode
import org.objectweb.asm.tree.InsnList
import org.objectweb.asm.tree.LdcInsnNode
import org.objectweb.asm.tree.MethodInsnNode
import org.objectweb.asm.tree.VarInsnNode
import java.io.File
import java.io.IOException
import java.io.InputStream
import java.io.OutputStream
import java.lang.management.ManagementFactory
import java.lang.management.ThreadMXBean
import java.net.URI
import java.util.Locale
import java.util.concurrent.CopyOnWriteArrayList
import java.util.concurrent.ExecutorService
import java.util.concurrent.Executors
import java.util.concurrent.ForkJoinPool
import java.util.concurrent.ForkJoinTask
import java.util.concurrent.Future
import java.util.concurrent.LinkedBlockingQueue
import java.util.concurrent.RecursiveTask
import java.util.concurrent.RejectedExecutionHandler
import java.util.concurrent.ThreadPoolExecutor
import java.util.concurrent.TimeUnit
import java.util.jar.JarFile
import java.util.zip.ZipEntry
import java.util.zip.ZipFile

class ConstEncryptTransformInvocation(
    private val delegate: TransformInvocation,
    private val transform: ConstEncryptTransform
) : TransformInvocation by delegate {
    private val NCPU = Runtime.getRuntime().availableProcessors()
    private val project = transform.project
    private val outputs = CopyOnWriteArrayList<File>()
    fun doIncrementalTransform() {
        doTransform(this::transformIncrementally)
    }

    private fun doTransform(block: (ExecutorService) -> Iterable<Future<*>>) {
        val executor = Executors.newFixedThreadPool(NCPU)
        try {
            block(executor).forEach {
                it.get()
            }
        } finally {
            executor.shutdown()
            executor.awaitTermination(1, TimeUnit.HOURS)
        }
    }

    @Suppress("NON_EXHAUSTIVE_WHEN")
    private fun doIncrementalTransform(jarInput: JarInput) {
        when (jarInput.status) {
            REMOVED -> jarInput.file.delete()
            CHANGED, ADDED -> {
                project.logger.info("Transforming ${jarInput.file}")
                outputProvider?.let { provider ->
                    jarInput.transform(
                        provider.getContentLocation(
                            jarInput.name,
                            jarInput.contentTypes,
                            jarInput.scopes,
                            Format.JAR
                        )
                    )
                }
            }
        }
    }

    private fun transformIncrementally(executor: ExecutorService) = this.inputs.map { input ->
        input.jarInputs.filter { it.status != NOTCHANGED }.map { jarInput ->
            executor.submit {
                doIncrementalTransform(jarInput)
            }
        } + input.directoryInputs.filter { it.changedFiles.isNotEmpty() }.map { dirInput ->
            val base = dirInput.file.toURI()
            executor.submit {
                doIncrementalTransform(dirInput, base)
            }
        }
    }.flatten()

    @Suppress("NON_EXHAUSTIVE_WHEN")
    private fun doIncrementalTransform(dirInput: DirectoryInput, base: URI) {
        dirInput.changedFiles.forEach { (file, status) ->
            when (status) {
                REMOVED -> {
                    outputProvider?.let { provider ->
                        provider.getContentLocation(
                            dirInput.name,
                            dirInput.contentTypes,
                            dirInput.scopes,
                            Format.DIRECTORY
                        ).parentFile.listFiles()?.asSequence()
                            ?.filter { it.isDirectory }
                            ?.map { File(it, dirInput.file.toURI().relativize(file.toURI()).path) }
                            ?.filter { it.exists() }
                            ?.forEach { it.delete() }
                    }
                    file.delete()
                }

                ADDED, CHANGED -> {
                    outputProvider?.let { provider ->
                        val root = provider.getContentLocation(
                            dirInput.name,
                            dirInput.contentTypes,
                            dirInput.scopes,
                            Format.DIRECTORY
                        )
                        val output = File(root, base.relativize(file.toURI()).path)
                        outputs += output
                        file.transform(output) { bytecode ->
                            bytecode.transform()
                        }
                    }
                }
            }
        }
    }

    fun doFullTransform() {
        this.outputs.clear()
        val executor = Executors.newFixedThreadPool(NCPU)
        try {
            transformFully(executor).forEach {
                it.get()
            }
        } finally {
            executor.shutdown()
            executor.awaitTermination(1, TimeUnit.HOURS)
        }
    }

    private fun QualifiedContent.transform(output: File) {
        outputs += output
        try {
            this.file.transform(output) { bytecode ->
                bytecode.transform()
            }
        } catch (e: Exception) {
            println("e===>${e.message}")
            e.printStackTrace()
        }
    }

    private fun ByteArray.transform(): ByteArray {
        return ClassWriter(ClassWriter.COMPUTE_MAXS).also { writer ->
            val classNode = ClassNode().also { klass ->
                ClassReader(this).accept(klass, 0)
            }

            classNode.methods.forEach { method ->
                val instructions = InsnList()
                method.instructions.forEach { insn ->
                    if (insn is MethodInsnNode &&
                        insn.owner == "com/ricky/constencrypt/api/ConstEncrypt" &&
                        insn.name == "wrap" &&
                        insn.desc == "(Ljava/lang/String;)Ljava/lang/String;"
                    ) {
                        val params = (insn.previous as? LdcInsnNode)?.cst as? String
                            ?: throw IllegalArgumentException("ConstEncrypt.wrap params is null")
                        val encryptParam = AESUtil.encrypt(params)
                        instructions.add(LdcInsnNode(encryptParam))
                        println("constencrypt class=${classNode.name} method=${method.name} ${params}->${encryptParam}")
                        instructions.add(
                            MethodInsnNode(
                                Opcodes.INVOKESTATIC, // 假设 decrypt 是静态方法
                                insn.owner,
                                "decrypt",
                                insn.desc,
                                false
                            )
                        )
                    } else {
                        // 保留原指令
                        instructions.add(insn)
                    }
                }

                // 更新方法的指令集
                method.instructions = instructions
            }

            // 将修改后的类写入 ClassWriter
            classNode.accept(writer)
        }.toByteArray()
    }

    private fun transformFully(executor: ExecutorService) = this.inputs.map {
        it.jarInputs + it.directoryInputs
    }.flatten().map { input ->
        executor.submit {
            val format = if (input is DirectoryInput) Format.DIRECTORY else Format.JAR
            outputProvider?.let { provider ->
                input.transform(
                    provider.getContentLocation(
                        input.name,
                        input.contentTypes,
                        input.scopes,
                        format
                    )
                )
            }
        }
    }

    class FileSearch internal constructor(
        private val roots: Iterable<File>,
        private val filter: (File) -> Boolean = { true }
    ) : RecursiveTask<Collection<File>>() {

        internal constructor(
            roots: Array<File>,
            filter: (File) -> Boolean = { true }
        ) : this(roots.toList(), filter)

        internal constructor(root: File, filter: (File) -> Boolean = { true }) : this(
            listOf(root),
            filter
        )

        override fun compute(): Collection<File> {
            val tasks = mutableListOf<RecursiveTask<Collection<File>>>()
            val result = mutableSetOf<File>()

            roots.forEach { root ->
                if (root.isDirectory) {
                    root.listFiles()?.let { files ->
                        FileSearch(files, filter).also { task ->
                            tasks.add(task)
                        }.fork()
                    }
                } else if (root.isFile) {
                    if (filter.invoke(root)) {
                        result.add(root)
                    }
                }
            }

            return result + tasks.flatMap { it.join() }
        }
    }

    private fun <T> ForkJoinTask<T>.execute(): T {
        val pool = ForkJoinPool()
        val result = pool.invoke(this)
        pool.shutdown()
        return result
    }

    @JvmOverloads
    fun File.search(filter: (File) -> Boolean = { true }): Collection<File> =
        FileSearch(this, filter).execute()

    private fun File.touch(): File {
        if (!this.exists()) {
            this.parentFile?.mkdirs()
            this.createNewFile()
        }
        return this
    }

    private fun ZipFile.transform(
        output: File,
        entryFactory: (ZipEntry) -> ZipArchiveEntry = ::ZipArchiveEntry,
        transformer: (ByteArray) -> ByteArray = { it -> it }
    ) = output.touch().outputStream().buffered().use {
        this.transform(it, entryFactory, transformer)
    }

    private fun InputStream.transform(transformer: (ByteArray) -> ByteArray): ByteArray {
        return transformer(readBytes())
    }

    private fun ZipFile.transform(
        output: OutputStream,
        entryFactory: (ZipEntry) -> ZipArchiveEntry = ::ZipArchiveEntry,
        transformer: (ByteArray) -> ByteArray = { it -> it }
    ) {
        val entries = mutableSetOf<String>()
        val creator = ParallelScatterZipCreator(
            ThreadPoolExecutor(
                NCPU,
                NCPU,
                0L,
                TimeUnit.MILLISECONDS,
                LinkedBlockingQueue<Runnable>(),
                Executors.defaultThreadFactory(),
                RejectedExecutionHandler { runnable, _ ->
                    runnable.run()
                })
        )
        //将jar包里的文件序列化输出
        entries().asSequence().forEach { entry ->
            if (!entries.contains(entry.name)) {
                val zae = entryFactory(entry)

                val stream = InputStreamSupplier {
                    when (entry.name.substringAfterLast('.', "")) {
                        "class" -> getInputStream(entry).use { src ->
                            try {
                                src.transform(transformer).inputStream()
                            } catch (e: Throwable) {
                                System.err.println("Broken class: ${this.name}!/${entry.name}")
                                getInputStream(entry)
                            }
                        }

                        else -> getInputStream(entry)
                    }
                }

                creator.addArchiveEntry(zae, stream)
                entries.add(entry.name)
            } else {
                System.err.println("Duplicated jar entry: ${this.name}!/${entry.name}")
            }
        }
        val zip = ZipArchiveOutputStream(output)
        zip.use { zipStream ->
            try {
                creator.writeTo(zipStream)
                zipStream.close()
            } catch (e: Exception) {
                zipStream.close()
                System.err.println("Duplicated jar entry: ${this.name}!")
            }
        }
    }

    private fun InputStream.redirect(file: File): Long =
        file.touch().outputStream().use { this.copyTo(it) }

    private fun ByteArray.redirect(file: File): Long = this.inputStream().use { it.redirect(file) }

    private fun File.transform(output: File, transformer: (ByteArray) -> ByteArray = { it -> it }) {
        when {
            isDirectory -> this.toURI().let { base ->
                this.search().parallelStream().forEach {
                    it.transform(File(output, base.relativize(it.toURI()).path), transformer)
                }
            }

            isFile -> when (extension.toLowerCase(Locale.ROOT)) {
                "jar" -> JarFile(this).use {
                    it.transform(output, { zipEntry -> JarArchiveEntry(zipEntry) }, transformer)
                }

                "class" -> this.inputStream().use {
                    it.transform(transformer).redirect(output)
                }

                else -> this.copyTo(output, true)
            }

            else -> throw IOException("Unexpected file: ${this.canonicalPath}")
        }
    }
}